import asyncio
import base64
import datetime
import hashlib  # 添加这个导入
import io
import json
import os
import subprocess
import tempfile
import threading
import time
import xml.etree.ElementTree as ET
from datetime import datetime, timedelta
from typing import Any

import jwt
import requests
from PIL import Image
from fastapi import APIRouter, HTTPException, Request
from qrcode.main import QRCode
from starlette.responses import HTMLResponse

from common.comm_utils import get_file_name
from controller.media_controller import generate_video_action
from server.audio_server import text2speech
from server.common_server import save_media_content
from server.image_server import create_image
from server.openai_server import create_server
from server.video_server import create_video_by_type, check_image_for_video, generate_action_for_video, \
    generate_3d_role_for_video
from common.cache import media_cache  # 改为从新模块导入
from server.wx_server import WxServer, verify_signature
from server.media_server import media_server

ai_server_instance = create_server()
router = APIRouter(prefix="/api/wx", tags=["微信相关"])

# 配置测试微信公众号信息
APP_ID = "wxa4e48ec87e373f87"
APP_SECRET = "607ac6d8167b519d1878f41322687bf7"
# 公众号配置
# APP_ID = "wx224a4add6e611d4c"
# APP_SECRET = "937902bff1226981a92b679fe6e50726"
# # 服务号配置
# APP_ID = "wx57da54bbf3042d6e"

ALLOWED_USERS = {
    "your_openid": {
        "nickname": "alive.cn",
        "openid": "oPAVM7PZQk00-zSiO_b1o4aixtAc"
    }
    # 可以添加更多允许的用户
}
JWT_SECRET = "your_secret_key"  # 使用与 auth_controller 相同的密钥

# 存储扫码状态和用户信息
scan_status = {}
user_info_cache = {}

# 添加 TOKEN 常量（这个需要与微信公众平台配置的Token一致）
TOKEN = "fucku"  # 替换成你在微信公众平台设置的Token

# 添加用户会话状态存储
user_sessions = {}

class UserSession:
    def __init__(self):
        self.last_text_time = None
        self.text_queue = []  # 最多存储10条消息
        # 修改图片存储逻辑，使用字典存储多个图片
        self.images = {
            'current': None,  # 最新上传的图片
            'base': None,     # 基础图片
            'mask': None      # 遮罩图片
        }
        self.current_video = None
        self.current_audio = None
        self.last_update_time = datetime.now()
        self.processed_msg_ids = set()  # 新增：用于存储已处理的消息ID
        self.waiting_for_materials = False  # 新增: 标记是否在等待素材
        self.required_materials = set()  # 新增: 记录需要的素材类型
    
    def update_text(self, text: str):
        self.last_text_time = datetime.now()
        self.text_queue.append(text)
        if len(self.text_queue) > 10:
            self.text_queue.pop(0)
        self.last_update_time = datetime.now()
    
    def update_media(self, media_type: str, media_url: str, image_role: str = 'current'):
        """
        更新媒体文件
        Args:
            media_type: 媒体类型 (video/image/audio)
            media_url: 媒体文件URL
            image_role: 图片角色 (current/base/mask)，仅当media_type为image时有效
        """
        if media_type == "video":
            self.current_video = media_url
        elif media_type == "image":
            self.images[image_role] = media_url
        elif media_type == "audio":
            self.current_audio = media_url
        self.last_update_time = datetime.now()
    
    def get_image(self, role: str = 'current') -> str:
        """
        获取指定角色的图片URL
        Args:
            role: 图片角色 (current/base/mask)
        Returns:
            str: 图片URL，如果不存在则返回None
        """
        return self.images.get(role)
    
    def should_process(self) -> bool:
        # 如果没有在等待素材,立即处理
        if not self.waiting_for_materials:
            return True
            
        # 如果在等待素材且超过60秒,则处理
        if self.last_text_time and (datetime.now() - self.last_text_time).total_seconds() >= 60:
            return True
            
        return False
    
    def should_cleanup(self) -> bool:
        return (datetime.now() - self.last_update_time).total_seconds() >= 1800
    
    def is_message_processed(self, msg_id: str) -> bool:
        return msg_id in self.processed_msg_ids
    
    def mark_message_processed(self, msg_id: str):
        self.processed_msg_ids.add(msg_id)
        # 保持集合大小在合理范围内
        if len(self.processed_msg_ids) > 1000:
            # 移除最旧的消息ID
            self.processed_msg_ids = set(list(self.processed_msg_ids)[-1000:])
    
    def reset_text_time(self):
        """重置最后文本时间，表示任务已处理完成"""
        self.last_text_time = None
    
    def set_waiting_materials(self, materials: set):
        """设置需要等待的素材类型"""
        self.waiting_for_materials = True
        self.required_materials = materials
        self.last_text_time = datetime.now()  # 重置等待时间
    
    def check_materials_ready(self) -> bool:
        """检查所需素材是否已准备好"""
        for material in self.required_materials:
            if material == 'image' and not self.get_image('current'):
                return False
            elif material == 'video' and not self.current_video:
                return False
            elif material == 'audio' and not self.current_audio:
                return False
        return True

# 新增：后台任务处理器
class SessionProcessor:
    def __init__(self):
        self.running = False
        self.thread = None
    
    def start(self):
        if not self.running:
            self.running = True
            self.thread = threading.Thread(target=self._process_loop)
            self.thread.daemon = True
            self.thread.start()
    
    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join()
    
    def _process_loop(self):
        while self.running:
            try:
                # 检查所有会话
                for user_id, session in list(user_sessions.items()):
                    # 只处理有未处理文本的会话
                    if session.last_text_time and session.should_process():
                        asyncio.run(process_session(user_id))
                # 清理过期会话
                cleanup_expired_sessions()
            except Exception as e:
                print(f"Session processing error: {e}")
            time.sleep(5)  # 每5秒检查一次

# 创建全局处理器实例
session_processor = SessionProcessor()

def create_token(user_info: dict):
    expiration = datetime.now() + timedelta(days=1)
    return jwt.encode(
        {
            "sub": user_info["openid"],
            "nickname": user_info["nickname"],
            "exp": expiration
        },
        JWT_SECRET,
        algorithm="HS256"
    )

@router.get("/qrcode")
async def get_qrcode():
    # 生成随机state用于防止CSRF攻击
    state = datetime.now().strftime("%Y%m%d%H%M%S")
    
    # 构建微信授权URL
    auth_url = (
        f"https://open.weixin.qq.com/connect/oauth2/authorize?"
        f"appid={APP_ID}&"
        f"redirect_uri=https://luckyzhou.online/api/wx/oauth_callback&"  # 需要替换成你的实际域名
        f"response_type=code&"
        f"scope=snsapi_userinfo&"
        f"state={state}#wechat_redirect"
    )
    
    # 生成包含授权URL的二维码
    qr = QRCode(version=1, box_size=10, border=5)
    qr.add_data(auth_url)
    qr.make(fit=True)
    
    # 将二维码转换为base64字符串
    img_buffer = io.BytesIO()
    qr.make_image(fill_color="black", back_color="white").save(img_buffer)
    qr_base64 = base64.b64encode(img_buffer.getvalue()).decode()
    
    return {"qrcode": f"data:image/png;base64,{qr_base64}", "state": state}

# 修改验证回调函数
@router.get("/callback")
async def wx_callback_verify(signature: str = "", timestamp: str = "", nonce: str = "", echostr: str = None):
    print(f"收到验证请求：signature={signature}, timestamp={timestamp}, nonce={nonce}, echostr={echostr}")
    
    if verify_signature(signature, timestamp, nonce, TOKEN):
        return HTMLResponse(content=echostr)
    return HTMLResponse(content="invalid signature")

def generate_reply_xml(msg_type: str, from_user: str, to_user: str, content: dict) -> str:
    """生成回复消息的XML"""
    try:
        return WxServer.get_instance(APP_ID, APP_SECRET).send_message(to_user, msg_type, content)
    except Exception as e:
        print(f"发送消息失败: {str(e)}")
        return "success"  # 即使发送失败也返回success避免微信服务器重试

def upload_media(media_type: str, file_content: bytes) -> str:
    """上传临时素材到微信服务器"""
    try:
        return WxServer.get_instance(APP_ID, APP_SECRET).upload_media(media_type, file_content)
    except Exception as e:
        print(f"上传媒体文件失败: {str(e)}")
        raise

def save_media_file(media_type: str, file_content: bytes, description: str = None, media_id: str = None) -> str:
    """保存媒体文件到本地目录"""
    return media_server.save_media_file(media_type, file_content, description, media_id)

@router.post("/callback")
async def wx_callback(request: Request):
    # 获取原始XML数据
    body = await request.body()
    print("body", body)

    try:
        # 解析XML
        xml_data = ET.fromstring(body)
        print("收到微信消息：", xml_data)
        
        # 提取消息ID和其他数据
        msg_id = xml_data.find('MsgId').text if xml_data.find('MsgId') is not None else None
        data = {
            'ToUserName': xml_data.find('ToUserName').text,
            'FromUserName': xml_data.find('FromUserName').text,
            'CreateTime': xml_data.find('CreateTime').text,
            'MsgType': xml_data.find('MsgType').text,
            'MsgId': msg_id,
        }

        # 初始化或获取用户会话
        from_user = data['FromUserName']
        if from_user not in user_sessions:
            user_sessions[from_user] = UserSession()
        session = user_sessions[from_user]
        
        # 获取或创建WxServer实例
        wx_server = WxServer.get_instance(APP_ID, APP_SECRET)
        
        try:
            # 检查消息是否已经处理过
            msg_id = data.get('MsgId')
            if msg_id and session.is_message_processed(msg_id):
                print(f"消息 {msg_id} 已经处理过，跳过处理")
                return HTMLResponse(content="success")
            
            print(f"开始处理新消息 {msg_id}")
            # 更新会话状态
            msg_type = data['MsgType']
            if msg_type == "text":
                session_processor.start()
                content = xml_data.find('Content').text
                session.update_text(content)
                
            elif msg_type == "image":
                media_id = xml_data.find('MediaId').text
                if not session.is_message_processed(media_id):
                    image_data = wx_server.download_media(media_id)
                    if image_data:
                        try:
                            # 尝试确定图片角色
                            image_role = determine_image_role(session)
                            local_url = save_media_file('image', image_data)
                            session.update_media("image", local_url, image_role)
                            
                            # 根据图片角色返回相应提示
                            if image_role == 'base':
                                reply_content = "已收到原图，请继续发送需要修改区域的遮罩图片"
                            elif image_role == 'mask':
                                reply_content = "已收到遮罩图片，正在处理您的修改请求"
                            else:
                                reply_content = "图片已收到"
                                
                            # 发送提示消息
                            return wx_server.send_message(
                                data['FromUserName'],
                                "text",
                                {"content": reply_content}
                            )
                            
                        except ValueError as ve:
                            # 处理特定的验证错误
                            return wx_server.send_message(
                                data['FromUserName'],
                                "text",
                                {"content": str(ve)}
                            )

            elif msg_type == "voice":
                media_id = xml_data.find('MediaId').text
                if not session.is_message_processed(media_id):
                    voice_data = wx_server.download_media(media_id)
                    if voice_data:
                        local_url = save_media_file('audio', voice_data)
                        session.update_media("audio", local_url)
                    
                    recognized_text = await wx_server.recognize_voice(media_id)
                    if recognized_text:
                        try:
                            intent = await recognize_intent([recognized_text], session)
                            if intent.get("operation") != "query":
                                # 如果识别出有效意图，将文本加入队列
                                session.update_text(recognized_text)
                                print(f"语音识别到意图: {intent}")
                            else:
                                print("语音未识别出特定意图，仅作为素材保存")
                        except Exception as e:
                            print(f"意图识别失败: {str(e)}")
            
            elif msg_type == "video":
                media_id = xml_data.find('MediaId').text
                if not session.is_message_processed(media_id):
                    video_data = wx_server.download_media(media_id)
                    if video_data:
                        local_url = save_media_file('video', video_data)
                        session.update_media("video", local_url)
            
            # 标记消息为已处理
            if msg_id:
                session.mark_message_processed(msg_id)
            
            return HTMLResponse(content="success")
            
        except Exception as e:
            print(f"处理消息失败: {str(e)}")
            return HTMLResponse(content="success")
            
    except Exception as e:
        print(f"解析XML失败: {str(e)}")
        return HTMLResponse(content="success")


async def process_session(user_id: str):
    session = user_sessions.get(user_id)
    if not session:
        return
    
    try:
        # 如果没有未处理的文本，直接返回
        if not session.text_queue or not session.last_text_time:
            return
            
        # 获取WxServer实例
        wx_server = WxServer.get_instance(APP_ID, APP_SECRET)
        
        # 识别意图
        intent = await recognize_intent(session.text_queue, session)
        
        # 如果在等待素材且素材未准备好,直接返回
        if session.waiting_for_materials and not session.check_materials_ready():
            return
            
        # 处理意图
        result = await process_intent(intent, session)
        
        # 只在发送失败时重试
        if result:
            max_retries = 3
            retry_count = 0
            
            while True:
                try:
                    # 使用WxServer发送消息
                    wx_server.send_message(user_id, result.get('type', 'text'), result)
                    print(f"消息发送成功")
                    # 重置会话状态，标记任务完成
                    session.reset_text_time()
                    session.text_queue.clear()  # 清空文本队列
                    session.waiting_for_materials = False  # 重置等待素材状态
                    session.required_materials.clear()  # 清空所需素材
                    break  # 成功发送后跳出循环
                    
                except Exception as e:
                    retry_count += 1
                    if retry_count >= max_retries:
                        print(f"发送消息失败，已达到最大重试次数: {str(e)}")
                        break
                    print(f"发送消息失败，正在重试 ({retry_count}/{max_retries}): {str(e)}")
                    await asyncio.sleep(1)

    except Exception as e:
        print(f"处理会话失败: {str(e)}")
        # 发生错误时也要重置会话状态，避免无限重试
        session.reset_text_time()
        session.text_queue.clear()
        session.waiting_for_materials = False
        session.required_materials.clear()

def cleanup_expired_sessions():
    expired_users = [
        user_id for user_id, session in user_sessions.items()
        if session.should_cleanup()
    ]
    for user_id in expired_users:
        user_sessions.pop(user_id, None)
        session_processor.stop()

@router.get("/check_scan")
async def check_scan(scene: str):
    print(f"检查扫码状态：scene={scene}, status={scan_status.get(scene)}")  # 添加日志

    # 检查是否已扫码
    if not scan_status.get(scene):
        return {"success": False, "message": "等待扫码"}
    
    # 获取用户信息
    user_info = user_info_cache.get(scene)
    if not user_info:
        raise HTTPException(status_code=500, detail="获取用户信息失败")
    
    # 检查用户是否已关注公众号
    if not user_info.get("subscribe"):
        raise HTTPException(status_code=403, detail="请先关注公众号")
    
    # 检查用户是否在允许列表中
    if user_info["openid"] not in ALLOWED_USERS:
        raise HTTPException(status_code=403, detail="未授权的用户")
    
    # 生成token
    token = create_token(user_info)
    
    # 清理缓存
    scan_status.pop(scene, None)
    user_info_cache.pop(scene, None)
    print("check_scan",{
        "success": True,
        "token": token,
        "user_info": {
            "nickname": user_info["nickname"],
            "openid": user_info["openid"],
            "avatar": user_info["headimgurl"]
        }
    })
    return {
        "success": True,
        "token": token,
        "user_info": {
            "nickname": user_info["nickname"],
            "openid": user_info["openid"],
            "avatar": user_info["headimgurl"]
        }
    }

# 可以添加一个清理过期场景值的定时任务
async def cleanup_expired_scenes():
    # 清理超过10分钟的场景值
    current_time = time.time()
    expired_scenes = [
        scene for scene, timestamp in scan_status.items()
        if current_time - timestamp > 600
    ]
    for scene in expired_scenes:
        scan_status.pop(scene, None)
        user_info_cache.pop(scene, None)

@router.get("/oauth_callback")
async def oauth_callback(code: str, state: str):
    # 使用code获取access_token
    token_url = (
        f"https://api.weixin.qq.com/sns/oauth2/access_token?"
        f"appid={APP_ID}&"
        f"secret={APP_SECRET}&"
        f"code={code}&"
        f"grant_type=authorization_code"
    )
    token_response = requests.get(token_url)
    token_data = token_response.json()
    
    if "errcode" in token_data:
        raise HTTPException(status_code=400, detail="获取access_token失败")
    
    # 使用access_token获取用户信息
    user_info_url = (
        f"https://api.weixin.qq.com/sns/userinfo?"
        f"access_token={token_data['access_token']}&"
        f"openid={token_data['openid']}&"
        f"lang=zh_CN"
    )
    user_response = requests.get(user_info_url)
    user_data = user_response.json()
    
    if "errcode" in user_data:
        raise HTTPException(status_code=400, detail="获取用户信息失败")
    
    # 检查用户是否在允许列表中
    if user_data["openid"] not in ALLOWED_USERS:
        return HTMLResponse(content="""
            <script>
                window.opener.postMessage({type: 'LOGIN_UNAUTHORIZED'}, '*');
                window.close();
            </script>
        """)
    
    # 生成token
    token = create_token({
        "openid": user_data["openid"],
        "nickname": user_data["nickname"]
    })
    
    # 返回HTML页面，通过postMessage将token和用户信息传递给opener窗口
    return HTMLResponse(content=f"""
        <script>
            window.opener.postMessage({{
                type: 'LOGIN_SUCCESS',
                data: {{
                    token: '{token}',
                    user_info: {{
                        nickname: '{user_data["nickname"]}',
                        openid: '{user_data["openid"]}',
                        avatar: '{user_data["headimgurl"]}'
                    }}
                }}
            }}, '*');
            window.close();
        </script>
    """)

async def recognize_intent(text_queue: list, session: UserSession) -> dict:
    # 分离当前问题和历史上下文
    current_query = text_queue[-1] if text_queue else ""
    history_context = "\n".join(text_queue[:-1]) if len(text_queue) > 1 else ""
    
    # 使用LLM判断是否为图片编辑意图
    edit_intent_prompt = f"""
    请判断以下用户输入是否表达了想要编辑/修改图片的意图。
    
    用户输入: {current_query}
    历史上下文: {history_context}
    
    请仅返回 true 或 false，不要包含任何其他内容。
    
    判断标准：
    1. 明确提到修改、编辑、更改图片
    2. 提到具体的图片编辑操作(如换背景、调色、涂鸦等)
    3. 表达希望改变图片某个部分或属性
    4. 提到图片风格转换
    5. 要求在图片上添加或删除内容
    """
    
    is_edit_response = await ai_server_instance.llm.ainvoke(edit_intent_prompt)
    is_edit_intent = is_edit_response.content.strip().lower() == "true"
    
    if is_edit_intent and not session.get_image('current'):
        # 如果是修改意图但没有图片，返回提示消息
        return {
            'operation': 'query',
            'parameters': {
                'text': "请先发送需要修改的图片"
            }
        }
    
    # 添加当前会话素材状态
    materials_status = f"""
    当前会话素材状态:
    - 图片: {"已上传: " + session.get_image('current') if session.get_image('current') else "未上传"}
    - 视频: {"已上传: " + session.current_video if session.current_video else "未上传"}  
    - 音频: {"已上传: " + session.current_audio if session.current_audio else "未上传"}
    """
    
    prompt = f"""
    基于以下用户当前问题、历史上下文和当前素材状态，识别用户的意图和所需的操作。请注意根据不同操作类型提供必要的参数：

    当前问题:
    {current_query}
    
    历史上下文:
    {history_context}
    
    {materials_status}
    
    可能的操作类型和所需参数：
    1. query - 普通对话查询
       - parameters: {{ "text": "用户查询内容" }}
       
    2. create_images - 生成或修改图片
       必要参数:
       - description: 图片描述
       可选参数:
       - negative_prompt: 负面提示词，默认为空
       - sketch_image_url: 参考图片URL (如果当前会话有上传图片)
       - base_image_url: 基础图片URL (用于图片修改)
       - mask_image_url: 遮罩图片URL (用于图片修改)
       - prompt_extend: 是否扩展提示词 (默认true)
       - watermark: 是否添加水印 (默认false)
       
    3. text_to_speech - 文字转语音
       必要参数:
       - text: 需要转换的文本内容
       
    4. generate_video_action - 提取视频动作
       必要参数:
       - video_url: 视频URL (需要当前会话有上传视频)
       
    5. create_action_video_by_video - 生成视频
       必要参数:
       - type: 视频类型，支持以下类型：
         ANIMATE (舞动人像):
           - image_url: 人像图片URL (必需)
           - template_id: 动作模板ID (必需，从视频中提取)
           可选参数:
           - use_ref_img_bg: 是否使用参考图片背景 (默认true)
           - video_ratio: 视频比例 (默认"9:16")
           
         EMO (悦动人像):
           - image_url: 人像图片URL (必需)
           - audio_id: 音频ID (必需)
           - face_bbox: 人脸边界框 (必需，从图片检测获得)
           - ext_bbox: 扩展边界框 (必需，从图片检测获得)
           可选参数:
           - style_level: 风格级别 (默认"normal")
           
         LIVE (灵动人像):
           - image_url: 虚拟肖像图片URL (必需)
           - audio_url: 音频URL (必需)
           可选参数:
           - template_id: 模板ID (默认"normal")
           - eye_move_freq: 眼睛移动频率 (默认0.5)
           - video_fps: 视频帧率 (默认30)
           - mouth_move_strength: 嘴部移动强度 (默认1)
           - paste_back: 是否粘贴背景 (默认"true")
           - head_move_strength: 头部移动强度 (默认0.7)
          
         VIDEO (声动人像):
           - video_url: 视频URL (必需)
           - audio_url: 音频URL (必需)
           可选参数:
           - video_extension: 视频扩展 (默认"false")
           
         Motion (幻影人像):
           - video_url: 视频URL (必需)
           - frame_index: 帧索引 (必需，从视频检测获得)
           - bbox: 边界框 (必需，从视频检测获得)
           - replacement_id: 替换ID列表 (必需，从3D角色生成获得)

    请根据用户消息内容和当前可用的素材，返回最合适的操作。如果所需素材不完整，应返回query类型并提示用户上传所需素材。
    
    请返回JSON格式的意图分析结果，包含：
    1. operation: 需要执行的操作
    2. type: 如果是create_action_video_by_video，指定具体type
    3. parameters: 根据操作类型提供所需的必要和可选参数
    
    注意：
    1. 禁止出现JSON格式以外的内容
    2. 确保提供的参数与当前会话状态匹配
    3. 对于需要多个素材的操作，确保所有必需素材都已上传
    4. 如果用户提供了具体参数值，优先使用用户提供的值
    5. 如果用户没有提供具体参数值，使用默认值
    """
    
    response = await ai_server_instance.llm.ainvoke(prompt)
    intent = json.loads(response.content)
    
    # 根据意图类型确定是否需要等待素材
    operation = intent.get('operation')
    if operation in ['create_images', 'text_to_speech', 'generate_video_action', 'create_action_video_by_video']:
        required_materials = set()
        
        if operation == 'create_images' and intent.get('parameters', {}).get('sketch_image_url'):
            required_materials.add('image')
            
        elif operation == 'generate_video_action':
            required_materials.add('video')
            
        elif operation == 'create_action_video_by_video':
            video_type = intent.get('type')
            if video_type == 'ANIMATE':
                required_materials.update(['image', 'video'])
            elif video_type in ['EMO', 'LIVE']:
                required_materials.update(['image', 'audio'])
            elif video_type == 'VIDEO':
                required_materials.update(['video', 'audio'])
            elif video_type == 'Motion':
                required_materials.update(['video', 'image'])
        
        # 如果需要素材且素材不完整,设置等待状态
        if required_materials and not session.check_materials_ready():
            session.set_waiting_materials(required_materials)
            # 返回提示用户上传所需素材的消息
            return {
                'operation': 'query',
                'parameters': {
                    'text': f"请上传以下素材: {', '.join(required_materials)}"
                }
            }
    
    return intent

async def process_intent(intent: dict, session: UserSession) -> dict:
    operation = intent.get('operation')
    parameters = intent.get('parameters', {})
    
    if operation == 'query':
        return await handle_query(parameters.get('text', session.text_queue[-1]))
        
    elif operation == 'create_images':
        # 补充图片生成所需的完整参数
        image_params = {
            'description': parameters.get('description', ' '.join(session.text_queue)),
            'negative_prompt': parameters.get('negative_prompt'),
            'sketch_image_url': session.get_image('current'),
            'prompt_extend': parameters.get('prompt_extend', True),
            'watermark': parameters.get('watermark', False)
        }
        
        # 如果存在base和mask图片，添加到extra_input
        if session.get_image('base') and session.get_image('mask'):
            image_params['extra_input'] = {
                'base_image_url': session.get_image('base'),
                'mask_image_url': session.get_image('mask')
            }
        
        # 调用create_image函数，传入所有参数
        result = create_image(
            description=image_params['description'],
            negative_prompt=image_params['negative_prompt'],
            sketch_image_url=image_params['sketch_image_url'],
            extra_input=image_params.get('extra_input'),
            prompt_extend=image_params['prompt_extend'],
        )
        
        # 处理返回结果
        if result and 'content' in result:
            image_data = base64.b64decode(result['content'])
            return {"type": "image", "content": image_data}
        return result
        
    elif operation == 'text_to_speech':
        return await handle_text_to_speech({
            'text': parameters.get('text', session.text_queue[-1])
        }, session)
        
    elif operation == 'generate_video_action':
        if not session.current_video:
            raise ValueError("No video provided")
        return await handle_generate_video_action({
            'video_url': session.current_video
        }, session)
        
    elif operation == 'create_action_video_by_video':
        video_type = intent.get('type', 'ANIMATE')
        # 根据不同视频类型准备参数
        video_params = {'type': video_type}
        
        if video_type == 'ANIMATE':
            if not session.get_image('current') or not session.current_video:
                raise ValueError("需要提供图片和视频才能生成舞动人像")
            video_params.update({
                'image_url': session.get_image('current'),
                'use_ref_img_bg': parameters.get('use_ref_img_bg', 'true'),
                'video_ratio': parameters.get('video_ratio', '9:16')
            })
            
        elif video_type == 'EMO':
            if not session.get_image('current') or not session.current_audio:
                raise ValueError("需要提供图片和音频才能生成悦动人像")
            video_params.update({
                'image_url': session.get_image('current'),
                'audio_id': session.current_audio,
                'style_level': parameters.get('style_level', 'normal')
            })
            
        elif video_type == 'LIVE':
            if not session.get_image('current') or not session.current_audio:
                raise ValueError("需要提供图片和音频才能生成灵动人像")
            video_params.update({
                'image_url': session.get_image('current'),
                'audio_url': session.current_audio,
                'template_id': parameters.get('template_id', 'normal'),
                'eye_move_freq': parameters.get('eye_move_freq', 0.5),
                'video_fps': parameters.get('video_fps', 30),
                'mouth_move_strength': parameters.get('mouth_move_strength', 1),
                'paste_back': parameters.get('paste_back', 'true'),
                'head_move_strength': parameters.get('head_move_strength', 0.7)
            })
            
        elif video_type == 'VIDEO':
            if not session.current_video or not session.current_audio:
                raise ValueError("需要提供视频和音频才能生成声动人像")
            video_params.update({
                'video_url': session.current_video,
                'audio_url': session.current_audio,
                'video_extension': parameters.get('video_extension', 'false')
            })
            
        elif video_type == 'Motion':
            if not session.current_video or not session.get_image('current'):
                raise ValueError("需要提供视频和图片才能生成幻影人像")
                
            # 检查视频是否合适
            check_result = check_image_for_video(type=video_type, video_url=session.current_video)
            if not check_result["output"]["check_pass"]:
                raise ValueError("提供的视频不适合生成幻影人像")
                
            # 生成3D角色
            ply_url = generate_3d_role_for_video(session.get_image('current'))
            if not ply_url:
                raise ValueError("无法生成3D角色模型")
                
            detected = check_result["output"]["detected"]
            frame_index = check_result["output"]["frame_index"]
            bbox = check_result["output"]["bbox"]
            
            video_params.update({
                'video_url': session.current_video,
                'detected': detected,
                'frame_index': frame_index,
                'bbox': bbox,
                'replacement_id': [ply_url]
            })
            
        else:
            raise ValueError(f"不支持的视频类型: {video_type}")

        # 调用视频生成服务
        result = create_video_by_type(type=video_type, json_post=video_params)
        
        if result and 'output' in result :
            if 'video_url' in result['output']:
                video_url = result['output']['video_url']
                # 下载视频
                video_data = requests.get(video_url).content
                # 返回视频数据
                return {"type": "video", "content": video_data}
            elif 'results' in result['results']:
                video_url = result['output']['results']['video_url']
                # 下载视频
                video_data = requests.get(video_url).content
                # 返回视频数据
                return {"type": "video", "content": video_data}
        
        raise ValueError("视频生成失败")
        
    raise ValueError(f"Unsupported operation: {operation}")

async def handle_query(query: str) -> dict:
    # 修改这里，确保我们调用 invoke() 方法
    response = await ai_server_instance.llm.ainvoke(query)
    return {"type": "text", "content": response.content}

async def handle_create_images(intent: dict, session: UserSession) -> dict:
    description = intent['parameters'].get('description', ' '.join(session.text_queue))
    result = create_image(
        description=description,
        sketch_image_url=session.get_image('current')
    )
    
    # 保存生成的图片到本地
    if result and 'content' in result:
        image_data = base64.b64decode(result['content'])
        return {"type": "image", "content": image_data}
    return result

async def handle_text_to_speech(intent: dict, session: UserSession) -> dict:
    text = intent['parameters'].get('text', session.text_queue[-1])
    audio_data = text2speech(text)
    
    # 保存生成的音频到本地
    local_url = save_media_file('audio', audio_data, text)
    print(f"Saved audio to {local_url}")
    return {"type": "audio", "content": audio_data}

async def handle_generate_video_action(intent: dict, session: UserSession) -> dict:
    if not session.current_video:
        raise ValueError("No video provided")
    
    response = await generate_video_action({
        "video_url": session.current_video
    })
    return {"type": "text", "content": response}

async def handle_create_action_video(intent: dict, session: UserSession) -> dict:
    """处理视频生成请求"""
    type = intent.get('type', 'ANIMATE')
    
    # 准备基础请求参数
    json_post = {"type": type}
    
    try:
        # 根据不同类型处理不同参数
        if type == 'ANIMATE':
            # 舞动人像需要图片和视频
            if not session.get_image('current') or not session.current_video:
                raise ValueError("需要提供图片和视频才能生成舞动人像")
                
            # 先检查图片是否合适
            check_result = check_image_for_video(type=type, image_url=session.get_image('current'))
            if not check_result["output"]["check_pass"]:
                raise ValueError("提供的图片不适合生成舞动人像")
                
            # 提取视频动作模板
            template_id = generate_action_for_video(session.current_video)
            if not template_id:
                raise ValueError("无法从视频中提取动作模板")
                
            json_post.update({
                "image_url": session.get_image('current'),
                "template_id": template_id
            })
            
        elif type in ['EMO', 'LIVE']:
            # 悦动/灵动人像需要图片和音频
            if not session.get_image('current') or not session.current_audio:
                raise ValueError("需要提供图片和音频才能生成悦动/灵动人像")
                
            # 检查图片是否合适
            check_result = check_image_for_video(type=type, image_url=session.get_image('current'))
            if not check_result["output"]["check_pass"]:
                raise ValueError(f"提供的图片不适合生成{type}人像")
                
            if type == 'EMO':
                # EMO需要额外的face_bbox和ext_bbox参数
                face_bbox = check_result["output"]["face_bbox"]
                ext_bbox = check_result["output"]["ext_bbox"]
                if not face_bbox or not ext_bbox:
                    raise ValueError("无法检测到人脸信息")
                    
                json_post.update({
                    "image_url": session.get_image('current'),
                    "audio_url": session.current_audio,
                    "face_bbox": face_bbox,
                    "ext_bbox": ext_bbox
                })
            else:  # LIVE
                json_post.update({
                    "image_url": session.get_image('current'),
                    "audio_url": session.current_audio
                })
                
        elif type == 'VIDEO':
            # 声动人像需要视频和音频
            if not session.current_video or not session.current_audio:
                raise ValueError("需要提供视频和音频才能生成声动人像")
                
            json_post.update({
                "video_url": session.current_video,
                "audio_url": session.current_audio
            })
            
        elif type == 'Motion':
            # 幻影人像需要视频和图片
            if not session.current_video or not session.get_image('current'):
                raise ValueError("需要提供视频和图片才能生成幻影人像")
                
            # 检查视频是否合适
            check_result = check_image_for_video(type=type, video_url=session.current_video)
            if not check_result["output"]["check_pass"]:
                raise ValueError("提供的视频不适合生成幻影人像")
                
            # 生成3D角色
            ply_url = generate_3d_role_for_video(session.get_image('current'))
            if not ply_url:
                raise ValueError("无法生成3D角色模型")
                
            detected = check_result["output"]["detected"]
            frame_index = check_result["output"]["frame_index"]
            bbox = check_result["output"]["bbox"]
            
            json_post.update({
                "video_url": session.current_video,
                "detected": detected,
                "frame_index": frame_index,
                "bbox": bbox,
                "replacement_id": [ply_url]
            })
            
        else:
            raise ValueError(f"不支持的视频类型: {type}")

        # 调用视频生成服务
        result = create_video_by_type(type=type, json_post=json_post)
        
        if result and 'output' in result :
            if 'video_url' in result['output']:
                video_url = result['output']['video_url']
                # 下载视频
                video_data = requests.get(video_url).content
                # 返回视频数据
                return {"type": "video", "content": video_data}
            elif 'results' in result['results']:
                video_url = result['output']['results']['video_url']
                # 下载视频
                video_data = requests.get(video_url).content
                # 返回视频数据
                return {"type": "video", "content": video_data}
        
        raise ValueError("视频生成失败")
        
    except Exception as e:
        print(f"生成视频失败: {str(e)}")
        return {"type": "text", "content": f"生成视频失败: {str(e)}"}

# 定期清理过期媒体文件的函数
def cleanup_expired_media():
    while True:
        try:
            media_server.cleanup_expired()
        except Exception as e:
            print(f"清理过期媒体文件失败: {str(e)}")
        time.sleep(3600)  # 每小时检查一次

# 启动清理线程
cleanup_thread = threading.Thread(target=cleanup_expired_media, daemon=True)
cleanup_thread.start()

def determine_image_role(session: UserSession) -> str:
    """
    根据会话上下文确定新上传图片的角色
    Returns:
        str: 图片角色 (current/base/mask)，或者提示信息
    """
    # 分析最近的文本消息来判断用户意图
    recent_texts = session.text_queue[-3:]  # 获取最近3条消息
    
    # 检查是否有图片修改相关关键词
    is_edit_intent = any(keyword in text for text in recent_texts 
                        for keyword in ['修改', '换背景', '涂鸦', '遮罩', 'mask'])
    
    if is_edit_intent:
        # 如果是图片修改意图，需要按顺序上传base和mask
        if not session.get_image('base'):
            # 如果还没有base图片，这张图片作为base
            return 'base'
        elif not session.get_image('mask'):
            # 如果已有base但没有mask，这张图片作为mask
            return 'mask'
        else:
            # 如果base和mask都有了，提示用户已经有完整的修改素材
            raise ValueError("已收到原图和遮罩图片，可以开始修改了")
    
    # 如果不是修改意图，作为普通图片处理
    return 'current'